#include "ratelimitconn.h"

RatelimitConn::RatelimitConn(Connection* pConn, bool splitable)
{
	real_conn_ = pConn;
	splitable_ = splitable;

	ProbingManager::Instance()->AddConnection(real_conn_);
	send_bucket_.AddListener(this);
}

RatelimitConn::~RatelimitConn()
{
	ProbingManager::Instance()->RemoveConnection(real_conn_);
	send_bucket_.RemoveListener(this);
}

void RatelimitConn::SetRateLimit(uint32 limit)//Bytes/sec
{
	if( limit == UINT_MAX ){
		send_bucket_.Init(NULL);
	}else{
		//rate_cfg_.rate_limit_ = limit;
		//rate_cfg_.rate_burst_ = limit;
		rate_cfg_.Init(limit, limit, 20);
		send_bucket_.Init(&rate_cfg_, false);
	}
}

int RatelimitConn::Send(byte* data, size_t size)
{
	//write according to the token bucket's limitation
	int total_bytes = 0;
	int total_cnt = 0;
	while(total_bytes<=size){
		int32 can_io_bytes = send_bucket_.GetLimit();
		if( !splitable_ && can_io_bytes < size ){
			//LOG4CPLUS_TRACE(g_logger," tokens not enough: "<<can_io_bytes<<" need send: "<<size);
			break;
		} 
		int written = real_conn_->Send(data+total_bytes, size-total_bytes);
		if( written > 0 ){
			int left = send_bucket_.RemoveTokens(written);
			total_bytes += written;
			++total_cnt;
		}else{
			//LOG4CPLUS_TRACE(g_logger,"written: "<<written<<" tokens: "<<send_bucket_.GetLimit());
			break;
		}
	}

	//if( total_bytes > 0 ){
	//	LOG4CPLUS_DEBUG(g_logger,"total sent bytes: "<< total_bytes<<" total_cnt: "<<total_cnt);
	//}

	return total_bytes;
}


//Set interval of probing, if set to INT_MAX, will disable active probing, but still can process probing packets
//Default 10 ms
void RatelimitConn::SetProbingInterval(int ms)
{
	real_conn_->SetProbingInterval(ms);
}

//if disabled, the connection will not process probing packets and will forward all data read to the app layer
void RatelimitConn::EnablePing(bool enabled)
{
	real_conn_->EnablePing(enabled);
}

IConnection::PROTO_TYPE RatelimitConn::GetProtoType()
{
	return real_conn_->GetProtoType();
}

const SocketAddress& RatelimitConn::RemoteAddress()
{
	return real_conn_->RemoteAddress();
}

void RatelimitConn::OnTokenFilled(TokenBucket * pBucket)
{
	LOG4CPLUS_DEBUG(g_logger,"token filled");
}

void RatelimitConn::OnTokenEmpty(TokenBucket * pBucket)
{
	LOG4CPLUS_DEBUG(g_logger,"token empty");
}

bool RatelimitConn::IsTimetoPing(uint32 now)
{
	return real_conn_->IsTimetoPing(now);
}

//The ping can be sent using APP data or Ping packets
//The default connection use Ping packets, but if the buffered connection, it can use ADD data, simply sent data
//out, and the underlying connection will decide automatically, if failed, should cope with the failure, and try
//ping when OnWrite Callbacked.
void RatelimitConn::DoPing()
{
	return real_conn_->DoPing();
}

bool RatelimitConn::Connected() const
{
	return real_conn_->Connected();
}
